module waterinfo
use prec
use iounits
use regioninfo
implicit none
   !----------------------------------------------------------------------------
   private
   public  water_init, add_water, lwater, H2O_init
   public  nH2Oatom, nH2Obond, nH2Oangle, H2OMass, H2OName
   public  watermolID, watertype, waterbond, waterangle, waterpos, waterchg
   public  HOBondCoeff, HOHAngCoeff, H2OPair
   !----------------------------------------------------------------------------
   logical :: lwater = .False.
   integer :: nwater = 0, nH2Oatom = 0, nH2Obond = 0, nH2Oangle = 0
   !----------------------------------------------------------------------------
   integer, allocatable :: watertype(:), waterbond(:,:), waterangle(:,:)
   integer, allocatable :: watermolID(:)
   real(q), allocatable :: waterpos(:,:), waterchg(:)
   !----------------------------------------------------------------------------
   integer :: H2Otype(3), H2Obond(2,2), H2OAngle(3,1)
   real(q) :: H2OPos(3,3), H2OChg(3), ang0 = 104.52D0/180.D0*3.141592653589793D0
   real(q) :: H2OMass(2) = (/ 15.9994D0, 1.00794D0 /)
   !----------------------------------------------------------------------------
   ! J. Chem Phys 125, 014702, 2006; J Chem Phys 121, 10096, 2004. TIP3P-Ew
   real(q), parameter :: HOBondCoeff(2) = (/ 450.D0, 0.9572D0  /)
   real(q), parameter :: HOHAngCoeff(2) = (/ 55.D0, 104.52D0   /)
   real(q), parameter :: H2OPair(2,2)   = reshape((/0.102D0,3.188D0,0.D0,0.D0/), (/2,2/))
   character(len=2)   :: H2OName(2) = (/ "O", "H" /)
   !----------------------------------------------------------------------------
contains
   !
   subroutine water_init()
   implicit none
      !-------------------------------------------------------------------------
      character(len=256) :: oneline
      integer            :: ioerr
      !-------------------------------------------------------------------------
      write(stdout, '(/,10x,"Would you like to insert water molecules into the system? (y/n)[n]: ", $)')
      read(stdin,'(A)',iostat=ioerr) oneline
      lwater = ioerr.eq.0.and.(oneline.eq.'y'.or.oneline.eq.'Y')
      !
      if (.not.lwater) return
      !
      ! initialize H2O molecule info
      H2Otype      = (/ 1, 2, 2 /)
      H2Obond(:,1) = (/ 1, 2 /)
      H2Obond(:,2) = (/ 1, 3 /)
      H2Oangle(:,1)= (/ 2, 1, 3 /)
      H2OChg       = (/ -0.830D0, 0.415D0, 0.415D0 /)
      H2OPos(:,1)  = (/ 0.D0, 0.D0, 0.D0 /)
      H2OPos(:,2)  = (/ 0.9572D0, 0.D0, 0.D0 /)
      H2OPos(:,3)  = (/ 0.9572D0*cos(ang0), 0.9572D0*sin(ang0), 0.D0 /)
      !
      ! ask for total number of water molecules
      write(stdout, '(10x,"How many water molecules would you like to insert? : ",$)')
      nwater = 0
      do while (ioerr.ne.0.or.nwater.lt.1)
         read(stdin,*,iostat=ioerr) nwater
      enddo
      nH2Oatom = nwater * 3
      nH2Obond = nwater * 2
      nH2Oangle= nwater
      !
      if ( allocated( watertype ) ) deallocate( watertype )
      if ( allocated( waterbond ) ) deallocate( waterbond )
      if ( allocated( waterangle) ) deallocate( waterangle)
      if ( allocated( waterpos  ) ) deallocate( waterpos  )
      if ( allocated( waterchg  ) ) deallocate( waterchg  )
      if ( allocated( waterMolID) ) deallocate( waterMolID)
      allocate( watertype(nH2Oatom), waterbond(2,nH2Obond), waterangle(3,nH2Oangle), &
      &         waterpos(3,nH2Oatom), waterchg(nH2Oatom), waterMolID(nH2Oatom) )
      !
      call init_region()
      !
   return
   end subroutine

   subroutine add_water()
   use cell, only: box, zboxlow, zboxhigh
   use wallinfo, only: LAddWalls
   implicit none
      !-------------------------------------------------------------------------
      integer :: i, istr, iend, iprev
      real(q) :: rndnum(6), TPI, rotatedH2O(3,3)
      logical :: loverlap
      !-------------------------------------------------------------------------
      TPI  = 8.D0 * atan(1.D0)
      !
      do i = 1, nwater
         loverlap = .true.
         istr     = (i-1) * 3 + 1
         iend     = i     * 3
         iprev    = istr  - 1
         ! To rotate and translate the molecule randomly, make sure that no atom will overlap.
         do while ( loverlap )
            call RANDOM_NUMBER( rndnum )
            rndnum(1:3) = rndnum(1:3) * TPI
            call rotate_molecular( rndnum(1), rndnum(2), rndnum(3), H2OPos, rotatedH2O, 3)
            call translate_molecular( rndnum(4), rndnum(5), rndnum(6), rotatedH2O, 3)
            if ( lregion ) loverlap = check_out_region(rotatedH2O)
            if (.not.loverlap.and.LAddWalls ) then
               loverlap = any(rotatedH2O(3,:).lt.zboxlow).or.any(rotatedH2O(3,:).gt.zboxhigh)
            endif
            if (.not.loverlap) loverlap = check_water_position(iprev, rotatedH2O)
         enddo
         ! position, atom type and chage of the i-th molecule
         waterpos(:, istr:iend)  = rotatedH2O
         watertype(istr:iend)    = H2Otype
         waterchg(istr:iend)     = H2OChg
         watermolID(istr:iend)   = i
         !
         ! bonds of the i-th molecule
         istr = (i-1) * 2 + 1
         iend = i     * 2
         waterbond(:,istr:iend) = H2Obond(:,:) + iprev
         !
         ! angles of the i-th molecule
         waterangle(:,i) = H2Oangle(:,1) + iprev
      enddo
   return
   end subroutine

   logical function check_water_position(iprev, pos)
   use cell, only: atpos, natom, box, rbox
   use wallinfo, only: LAddWalls, WallPos, nWallAtom, WallRemain
   implicit none
      !-------------------------------------------------------------------------
      integer, intent(in) :: iprev
      real(q), intent(in) :: pos(3,3)
      !-------------------------------------------------------------------------
      real(q) :: Rmin2(3) = (/ 3.15D0, 1.33D0, 1.33D0 /), Rwmin = 7.0D0
      !-------------------------------------------------------------------------
      integer :: i, j, k
      real(q) :: R(3), dist
      !-------------------------------------------------------------------------
      check_water_position = .false.
      !
      ! all liquid atoms
liquid:do i = 1, natom
         do j = 1, 3
            R = pos(:,j) - atpos(:,i)
            R = R - nint(R*rbox)*box
            dist = R(1)*R(1) + R(2)*R(2) + R(3)*R(3)
            check_water_position = dist.lt.Rmin2(j)
            if (check_water_position) exit liquid
         enddo
      enddo liquid
      ! Wall atoms
      if ( LAddWalls.and.(.not.check_water_position) ) then
wall:    do i = 1, nWallAtom
            if ( WallRemain(i).eq.0 ) cycle
            do j = 1, 3
               R = pos(:,j) - WallPos(:,i)
               R(1:2) = R(1:2) - nint( R(1:2)*rbox(1:2) ) * box(1:2)
               dist   = R(1)*R(1) + R(2)*R(2) + R(3)*R(3)
               !
               check_water_position = dist.lt.Rwmin
               if (check_water_position) exit wall
            enddo
         enddo wall
      endif
      ! Previous water atoms
water:do i = 1, iprev
         do j = 1, 3
            R = pos(:,j) - waterpos(:, i)
            R = R - nint( R*rbox ) * box
            dist = R(1)*R(1) + R(2)*R(2) + R(3)*R(3)
            check_water_position = dist.lt.Rmin2(j)
            if (check_water_position) exit water
         enddo
      enddo water
   return
   end function check_water_position

   subroutine H2O_Init()
   use cell
   use general
   implicit none
      !----------------------------------------------------------------------------
      integer  :: i, moltype, fftype = 1
      character(len=256) :: oneline
      !-------------------------------------------------------------------------
      write(stdout, '(/,10x,"Which kind of force field will be used to describe water? ")')
      write(stdout, '(12x, "1. TIP3P cutoff;")')
      write(stdout, '(12x, "2. TIP3P with long range;")')
      write(stdout, '(12x, "3. TIP4P cutoff;")')
      write(stdout, '(12x, "4. TIP4P ice;")')
      write(stdout, '(12x, "5. TIP4P 2005;")')
      write(stdout, '(12x, "6. TIP4P with long range;")')
      write(stdout, '(12x, "7. SPC cut at 9 Ang;")')
      write(stdout, '(12x, "8. SPC/E cut at 9 Ang;")')
      write(stdout, '(10x, "Your choice [1]: ", $)')
      read(stdin, '(A)', iostat=ioerr) oneline
      if (ioerr.eq.0) then
         oneline = trim(oneline) // " 1"
         read(oneline, *, iostat=ioerr) fftype
         if (ioerr.ne.0) fftype = 1
      endif
      fftype = MAX(1, MIN(fftype, 6))
      !----------------------------------------------------------------------------
      subnamep = subname
      subname  = 'H2O_init'
      !
      if ( allocated( bondsMol  ) ) deallocate( bondsMol  )
      if ( allocated( anglsMol  ) ) deallocate( anglsMol  )
      if ( allocated( dihrsMol  ) ) deallocate( dihrsMol  )
      if ( allocated( attypMol  ) ) deallocate( attypMol  )
      if ( allocated( atposMol  ) ) deallocate( atposMol  )
      if ( allocated( molIDMol  ) ) deallocate( molIDMol  )
      if ( allocated( rotatedMol) ) deallocate( rotatedMol)
      if ( allocated( chargeMol ) ) deallocate( chargeMol )
      if ( allocated( EMass     ) ) deallocate( EMass     )
      if ( allocated( sigma     ) ) deallocate( sigma     )
      if ( allocated( epison    ) ) deallocate( epison    )
      if ( allocated( bondcoef  ) ) deallocate( bondcoef  )
      if ( allocated( anglcoef  ) ) deallocate( anglcoef  )
      if ( allocated( torscoef  ) ) deallocate( torscoef  )
      if ( allocated( DistMin   ) ) deallocate( DistMin   )
      if ( allocated( EName     ) ) deallocate( EName     )
      !
      ntypeMol   = 2  ! 1: O; 2: H
      nbndtypMol = 1
      nangtypMol = 1
      ndihtypMol = 0
      !
      allocate( EMass(ntypeMol), sigma(ntypeMol), epison(ntypeMol),           &
      &         bondcoef(2, nbndtypMol), anglcoef(2, nangtypMol),             &
      &         DistMin(ntypeMol), EName(ntypeMol+1) )
      DistMin  = (/ 2.2D0, 1.2D0 /)
      EName    = (/ 'O', 'H' /)
      !
      molform    = 'H2O'
      den0       = 1.D0
      natomMol   = 3
      nbondsMol  = 2
      nanglsMol  = 1
      ndihrsMol  = 0
      !
      allocate( bondsMol(3, nbondsMol), AnglsMol(4, nanglsMol),        &
      &         attypMol(natomMol),                                    &
      &         atposMol(3,natomMol), chargeMol(natomMol), molIDMol(natomMol) )                
      !
      attypMol = (/ 1, 2, 2 /)
      !
      atposMol(:, 1) = (/ 0.D0, 0.D0, 0.D0 /)
      atposMol(:, 2) = (/ 0.9572D0, 0.D0, 0.D0 /)
      atposMol(:, 3) = (/ 0.9572D0*cos(104.52D0/180.D0*3.141592653589793D0),&
      &              0.9572D0*sin(104.52D0/180.D0*3.141592653589793D0), 0.D0 /)
      !
      !                    Type, atom 1, atom 2
      bondsMol(:, 1) = (/   1,  1,  2 /)
      bondsMol(:, 2) = (/   1,  1,  3 /)
      !                   Type,  atom1,  atom2, atom3
      AnglsMol(:, 1) = (/  1,  2,  1,  3 /)
      !
      molIDMol  = (/ 1, 1, 1 /)
      !
      ! To define the force field parameters
      select case (fftype)
      case (2) ! TIP3P long range
         !
         molname = "water-tip3p-long"
         write(stdout, '(10x, "With TIP3P long range: fix-shake can be used.")')
         EMass    = (/ 15.9994D0, 1.008D0 /)
         chargeMol = (/ -0.830D0, 0.415D0, 0.415D0 /)
         sigma    = (/ 3.188D0, 0.0D0 /)
         epison   = (/ 0.102D0, 0.0D0 /)
         bondcoef = reshape( (/ 450.D0, 0.9572D0 /),(/2, nbndtypMol/) ) ! O-H
         anglcoef = reshape( (/ 55.D0, 104.52D0  /),(/2, nangtypMol/) ) ! H-O-H
         !
      case (3) ! TIP4P cut
         !
         molname = "water-tip4p-cut"
         write(stdout, '(10x, "With TIP4P cut: OM dist 0.15, cutoff at 8.5 with fix-shake.")')
         EMass    = (/ 15.9994D0, 1.008D0 /)
         chargeMol = (/ -1.04D0, 0.52D0, 0.52D0 /)
         sigma    = (/ 3.1536D0, 0.0D0 /)
         epison   = (/ 0.1550D0, 0.0D0 /)
         bondcoef = reshape( (/ 450.D0, 0.9572D0 /),(/2, nbndtypMol/) ) ! O-H
         anglcoef = reshape( (/ 55.D0, 104.52D0  /),(/2, nangtypMol/) ) ! H-O-H
         !
      case (4) ! TIP4P ice
         !
         molname = "water-tip4p-cut-ice"
         write(stdout, '(10x, "With TIP4P/ice: OM dist 0.1577, cutoff at 8.5 with fix-shake.")')
         EMass    = (/ 15.9994D0, 1.008D0 /)
         chargeMol = (/ -1.1794D0, 0.5897D0, 0.5897D0 /)
         sigma    = (/ 3.1668D0,  0.0D0 /)
         epison   = (/ 0.21084D0, 0.0D0 /)
         bondcoef = reshape( (/ 450.D0, 0.9572D0 /),(/2, nbndtypMol/) ) ! O-H
         anglcoef = reshape( (/ 55.D0, 104.52D0  /),(/2, nangtypMol/) ) ! H-O-H
         !
      case (5) ! TIP4P 2005
         !
         molname = "water-tip4p-cut-2005"
         write(stdout, '(10x, "With TIP4P 2005: OM dist 0.1546, cutoff at 8.5 with fix-shake.")')
         EMass    = (/ 15.9994D0, 1.008D0 /)
         chargeMol = (/ -1.1128D0, 0.5564D0, 0.5564D0 /)
         sigma    = (/ 3.1589D0, 0.0D0 /)
         epison   = (/ 0.1852D0, 0.0D0 /)
         bondcoef = reshape( (/ 450.D0, 0.9572D0 /),(/2, nbndtypMol/) ) ! O-H
         anglcoef = reshape( (/ 55.D0, 104.52D0  /),(/2, nangtypMol/) ) ! H-O-H
         !
      case (6) ! TIP4P long range
         !
         molname = "water-tip4p-long"
         write(stdout, '(10x, "With TIP4P long range: OM dist 0.125 with fix-shake.")')
         EMass    = (/ 15.9994D0, 1.008D0 /)
         chargeMol = (/ -1.0484D0, 0.5242D0, 0.5242D0 /)
         sigma    = (/ 3.16435D0, 0.0D0 /)
         epison   = (/ 0.16275D0, 0.0D0 /)
         bondcoef = reshape( (/ 450.D0, 0.9572D0 /),(/2, nbndtypMol/) ) ! O-H
         anglcoef = reshape( (/ 55.D0, 104.52D0  /),(/2, nangtypMol/) ) ! H-O-H
         !
      case (7) ! SPC model
         !
         molname = "water-SPC"
         write(stdout, '(10x, "With SPC model, cutoff at 9 with fix-shake.")')
         EMass    = (/ 15.9994D0, 1.008D0 /)
         chargeMol = (/ -0.82D0, 0.41D0, 0.41D0 /)
         sigma    = (/ 3.166D0,  0.0D0 /)
         epison   = (/ 0.1553D0, 0.0D0 /)
         bondcoef = reshape( (/ 450.D0, 1.0D0 /),(/2, nbndtypMol/) ) ! O-H
         anglcoef = reshape( (/ 55.D0, 109.47D0  /),(/2, nangtypMol/) ) ! H-O-H
         !
      case (8) ! SPC/E model
         !
         molname = "water-SPC/E"
         write(stdout, '(10x, "With SPC model, cutoff at 9 with fix-shake.")')
         EMass    = (/ 15.9994D0, 1.008D0 /)
         chargeMol = (/ -0.8476D0, 0.4238D0, 0.4238D0 /)
         sigma    = (/ 3.166D0,  0.0D0 /)
         epison   = (/ 0.1553D0, 0.0D0 /)
         bondcoef = reshape( (/ 450.D0, 1.0D0 /),(/2, nbndtypMol/) ) ! O-H
         anglcoef = reshape( (/ 55.D0, 109.47D0  /),(/2, nangtypMol/) ) ! H-O-H
         !
      case default ! TIP3P cutoff
         !
         molname = "water-tip3p-cut"
         write(stdout, '(10x, "With TIP3P cutoff: fix-shake can be used.")')
         EMass    = (/ 15.9994D0, 1.008D0 /)
         chargeMol = (/ -0.8340D0, 0.417D0, 0.417D0 /)
         sigma    = (/ 3.1507D0, 0.400D0 /)
         epison   = (/ 0.1521D0, 0.046D0 /)
         bondcoef = reshape( (/ 450.D0, 0.9572D0 /),(/2, nbndtypMol/) ) ! O-H
         anglcoef = reshape( (/ 55.D0, 104.52D0  /),(/2, nangtypMol/) ) ! H-O-H
         !
      end select

      nMol      = 1
      molmass   = 0.D0
      do i = 1, natomMol
         molmass = molmass + EMass( attypMol(i) )
      enddo
      forall ( i=1:3 ) minbox(i) = maxval(atposMol(i,:)) - minval(atposMol(i,:))
      minlen = maxval( minbox )
      allocate( rotatedMol( 3, natomMol) )
      !
      subname = subnamep
   return
   end subroutine
end module
